import re
import arrow

from app.libs.redprint import RedPrint
from app.db.models.vulndb import KbVuln
from app.cache import cache
from app.libs.score import Score
from app.libs.utility import make_cache_key_global
from app.libs.utility import api_v3_response, get_request_args_v3


api = RedPrint('vuln')


@api.route('/vuln-type', methods=['GET'])
@cache.cached(timeout=600, make_cache_key=make_cache_key_global)
def get_vuln_type():
    query = {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False}
    data = KbVuln._get_collection().distinct("vuln_type.zh_cn", query)
    return api_v3_response(data=data)


@api.route('/vuln-list', methods=['GET'])
@cache.cached(timeout=600, make_cache_key=make_cache_key_global)
def get_vuln_info():
    request_args = get_request_args_v3(["vuln_name", "vuln_type", "update_start_at", "update_end_at", "page", "limit"])

    query = {'source': {'$in': ['kscan', 'nscan', 'cve', 'nuclei']}, 'is_deleted': False}
    if vuln_name := request_args.get("vuln_name"):
        query['name.zh_cn'] = {'$regex': vuln_name.replace('(', r'\(').replace(')', r'\)')}
    if vuln_type := request_args.get("vuln_type"):
        query['vuln_type.zh_cn'] = {'$regex': vuln_type}
    update_start = request_args.get("update_start_at")
    update_end = request_args.get("update_end_at")
    if update_start:
        try:
            update_start = arrow.get(update_start).datetime
        except arrow.parser.ParserError:
            return api_v3_response(status="fail", errcode=3040001, data={})
    if update_end:
        try:
            update_end = arrow.get(update_end).datetime
        except arrow.parser.ParserError:
            return api_v3_response(status="fail", errcode=3040001, data={})
    if update_start:
        query.setdefault('updated_at', {})['$gte'] = update_start
    if update_end:
        query.setdefault('updated_at', {})['$lte'] = update_end

    page, limit = request_args.get("page", 1), request_args.get("limit", 10)
    data = KbVuln.objects.find(query)
    total = data.count()

    if (page - 1) * limit >= total:
        results = {"page": page, "limit": limit, "count": total, "results": []}
        return api_v3_response(data=results)

    data = data.order_by("-updated_at").paginate(page, limit)
    severity_map = Score.severityMap
    new_data = []
    for item in data.items:
        new_data.append({
            'vuln_name': item.name.zh_cn,
            'vuln_type': item.vuln_type[0].zh_cn if item.vuln_type else '',
            'level': severity_map.get(item.severity, {}).get("zh_CN", "低危"),
            'published_at': arrow.get(item.published_at).isoformat() if item.published_at else '',
            'updated_at': arrow.get(item.updated_at).isoformat() if item.updated_at else '',
            'cnnvd': item.cnnvd,
            'cve': item.cve,
            'owasp': item.owasp,
            'description': item.description.zh_cn,
            'impact': item.impact.zh_cn if item.impact else '',
            'recommendation': item.recommendation.zh_cn if item.recommendation else '',
            'reference': item.reference,
            'zoomeye_dork': f'根据ZoomEye网络空间搜索引擎关键字 {item.zoomeye_dork} 对潜在可能目标进行搜索，共得到 {item.zoomeye_count} 条IP历史记录。',
        })
    results = {"page": page, "limit": limit, "count": total, "results": new_data}
    return api_v3_response(data=results)
